import os

import joblib
import pandas as pd
import numpy as np
from utils.data_load import dataload
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import GridSearchCV, StratifiedKFold
from sklearn.metrics import classification_report, confusion_matrix, roc_auc_score
from xgboost import XGBClassifier
from category_encoders import TargetEncoder
from imblearn.over_sampling import SMOTE

# ======================================================
# 1️⃣ 数据加载
# ======================================================
X_train, X_test, y_train, y_test = dataload()

# ======================================================
# 2️⃣ 类别特征目标编码
# ======================================================
categorical_features = ['BusinessTravel', 'Department', 'EducationField', 'JobRole', 'MaritalStatus', 'OverTime']

encoder = TargetEncoder(cols=categorical_features)
X_train[categorical_features] = encoder.fit_transform(X_train[categorical_features], y_train)
X_test[categorical_features] = encoder.transform(X_test[categorical_features])

# ======================================================
# 3️⃣ 特征工程：比例 & log 变换
# ======================================================
X_train['YearsAtCompany_Age_ratio'] = X_train['YearsAtCompany'] / (X_train['Age'] + 1e-5)
X_test['YearsAtCompany_Age_ratio'] = X_test['YearsAtCompany'] / (X_test['Age'] + 1e-5)

X_train['TotalWorkingYears_Age_ratio'] = X_train['TotalWorkingYears'] / (X_train['Age'] + 1e-5)
X_test['TotalWorkingYears_Age_ratio'] = X_test['TotalWorkingYears'] / (X_test['Age'] + 1e-5)

X_train['CurrentRole_Company_ratio'] = X_train['YearsInCurrentRole'] / (X_train['YearsAtCompany'] + 1e-5)
X_test['CurrentRole_Company_ratio'] = X_test['YearsInCurrentRole'] / (X_test['YearsAtCompany'] + 1e-5)

for col in ['MonthlyIncome', 'DistanceFromHome', 'NumCompaniesWorked']:
    if col in X_train.columns:
        X_train[col] = np.log1p(X_train[col])
        X_test[col] = np.log1p(X_test[col])

categorical_features = X_train.select_dtypes(include=['object']).columns.tolist()

# 创建目标编码器
encoder = TargetEncoder(cols=categorical_features, smoothing=0.3)
X_train[categorical_features] = encoder.fit_transform(X_train[categorical_features], y_train)
X_test[categorical_features] = encoder.transform(X_test[categorical_features])

# ======================================================
# 4️⃣ 样本平衡（SMOTE）
# ======================================================
sm = SMOTE(random_state=42)
X_train_res, y_train_res = sm.fit_resample(X_train, y_train)

# ======================================================
# 5️⃣ 标准化
# ======================================================
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train_res)
X_test_scaled = scaler.transform(X_test)

# ======================================================
# 6️⃣ GridSearchCV 超参数调优
# ======================================================
param_grid = {
    'n_estimators': [200, 400, 600],
    'max_depth': [3, 5, 7],
    'learning_rate': [0.01, 0.05, 0.1],
    'subsample': [0.7, 0.8, 1.0],
    'colsample_bytree': [0.7, 0.8, 1.0],
}

base_model = XGBClassifier(
    random_state=42,
    # use_label_encoder=False,
    eval_metric='logloss',
    scale_pos_weight=(len(y_train_res)-sum(y_train_res))/sum(y_train_res)
)

cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)

grid_search = GridSearchCV(
    estimator=base_model,
    param_grid=param_grid,
    scoring='roc_auc',   # 以 AUC 为评估指标
    cv=cv,
    n_jobs=-1,
    verbose=2
)

print("🔍 正在进行超参数搜索...")
grid_search.fit(X_train_scaled, y_train_res)

print("\n===== ✅ 最优参数 =====")
print(grid_search.best_params_)

print("\n===== 🌟 最优CV得分 =====")
print(grid_search.best_score_)

# ======================================================
# 7️⃣ 使用最优模型在测试集上评估
# ======================================================
best_model = grid_search.best_estimator_
os.makedirs("../model", exist_ok=True)
joblib.dump(best_model, "../model/xgboost_model.pkl")
y_pred = best_model.predict(X_test_scaled)
y_proba = best_model.predict_proba(X_test_scaled)[:, 1]

print("\n===== 分类报告 =====")
print(classification_report(y_test, y_pred))
print("\n===== 混淆矩阵 =====")
print(confusion_matrix(y_test, y_pred))
print("\n===== ROC AUC =====")
print(roc_auc_score(y_test, y_proba))
